package com.yjj.demo.common.interceptor;

import com.yjj.demo.common.annotation.ApiLimiter;
import lombok.extern.log4j.Log4j2;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.TimeUnit;

/**
 * 简单限速器
 *
 * @author yangjiajia
 * @createdAt 2023/1/18 15:33
 */
@Log4j2
public class ApiLimiterInterceptor implements HandlerInterceptor {

    @Autowired
    private RedisTemplate<String, String> redisTemplate;

    private static final String CONTENT_APPLICATION_JSON_UTF8 = "application/json;charset=UTF-8";

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            ApiLimiter apiLimiter = handlerMethod.getMethodAnnotation(ApiLimiter.class);
            if (null == apiLimiter) {
                return true;
            }
            String userId = request.getParameter("userId");
            String key = String.format("api_limiter_%s", userId);
            Boolean exist = redisTemplate.opsForValue().setIfAbsent(key, String.valueOf(1), apiLimiter.time(), TimeUnit.SECONDS);
            if (Boolean.FALSE.equals(exist)) {
                log.info("url:{},请求频率过高,userId:{}", request.getRequestURI(), userId);
                response.setContentType(CONTENT_APPLICATION_JSON_UTF8);
                response.getWriter().print(apiLimiter.msg());
                return false;
            }
        }
        return true;
    }
}
